#!/usr/bin/env python3
"""
MCP工具演示脚本

演示如何使用MySQL MCP服务器的各种工具功能。
"""

import asyncio
import json
import sys
import os
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from src.mysql_mcp_server.config import DatabaseConfig
from src.mysql_mcp_server.mcp_server import create_mcp_server


async def demo_mcp_tools():
    """演示MCP工具的使用"""
    print("🚀 MySQL MCP服务器工具演示")
    print("=" * 50)
    
    try:
        # 加载配置
        config_file = project_root / ".env"
        config = DatabaseConfig.from_env(str(config_file))
        print(f"✅ 配置加载成功: {config.host}:{config.port}/{config.database}")
        
        # 创建MCP服务器
        mcp_server = await create_mcp_server(config)
        print("✅ MCP服务器创建成功")
        
        # 启动服务器（初始化数据库连接池）
        await mcp_server.start()
        print("✅ MCP服务器启动完成")
        
        # 初始化MCP协议（设置initialized标志）
        initialize_request = {
            "jsonrpc": "2.0",
            "id": 0,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "clientInfo": {
                    "name": "demo-client",
                    "version": "1.0.0"
                }
            }
        }
        
        response = await mcp_server.handle_request(json.dumps(initialize_request))
        init_result = json.loads(response)
        print("✅ MCP协议初始化完成")
        
        # 演示1: 列出所有工具
        print("\n📋 可用的MCP工具:")
        print("-" * 30)
        
        list_tools_request = {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'tools/list',
            'params': {}
        }
        
        response = await mcp_server.handle_request(json.dumps(list_tools_request))
        tools_result = json.loads(response)
        
        if tools_result and 'result' in tools_result and tools_result['result'] and 'tools' in tools_result['result']:
            for tool in tools_result['result']['tools']:
                print(f"🔧 {tool['name']}: {tool['description']}")
        else:
            print("❌ 无法获取工具列表")
        
        # 演示2: 列出数据库表
        print("\n📊 数据库表列表:")
        print("-" * 30)
        
        list_tables_request = {
            'jsonrpc': '2.0',
            'id': 2,
            'method': 'tools/call',
            'params': {
                'name': 'list_tables',
                'arguments': {}
            }
        }
        
        response = await mcp_server.handle_request(json.dumps(list_tables_request))
        tables_result = json.loads(response)
        
        if tables_result and 'result' in tables_result and tables_result['result'] and 'content' in tables_result['result']:
            tables_info = json.loads(tables_result['result']['content'][0]['text'])
            if tables_info and 'tables' in tables_info:
                print(f"找到 {tables_info['count']} 个表:")
                for table_name in tables_info['tables']:
                    print(f"📋 {table_name}")
            else:
                print("❌ 无法获取表信息")
        else:
            print("❌ 无法获取表列表")
        
        # 演示3: 执行简单查询
        print("\n🔍 执行SQL查询演示:")
        print("-" * 30)
        
        query_request = {
            'jsonrpc': '2.0',
            'id': 3,
            'method': 'tools/call',
            'params': {
                'name': 'execute_query',
                'arguments': {
                    'query': 'SELECT COUNT(*) as table_count FROM information_schema.tables WHERE table_schema = DATABASE()',
                    'limit': 1
                }
            }
        }
        
        response = await mcp_server.handle_request(json.dumps(query_request))
        query_result = json.loads(response)
        
        if query_result and 'result' in query_result and query_result['result'] and 'content' in query_result['result']:
            query_info = json.loads(query_result['result']['content'][0]['text'])
            if query_info and 'results' in query_info:
                print(f"✅ 查询成功，返回 {len(query_info['results'])} 行数据:")
                for row in query_info['results']:
                    print(f"   {row}")
            else:
                print("❌ 无法获取查询结果")
        else:
            print("❌ 查询执行失败")
        
        # 演示表结构查询
        if 'tables_info' in locals() and tables_info and 'tables' in tables_info and tables_info['tables']:
            first_table = tables_info['tables'][0]
            print(f"\n🏗️  表结构演示 - {first_table}:")
            print("-" * 30)
            
            describe_request = {
                'jsonrpc': '2.0',
                'id': 4,
                'method': 'tools/call',
                'params': {
                    'name': 'describe_table',
                    'arguments': {
                        'table_name': first_table
                    }
                }
            }
            
            response = await mcp_server.handle_request(json.dumps(describe_request))
            describe_result = json.loads(response)
            
            if describe_result and 'result' in describe_result and describe_result['result'] and 'content' in describe_result['result']:
                # 获取文本内容并解析JSON
                text_content = describe_result['result']['content'][0]['text']
                describe_info = json.loads(text_content)
                
                if describe_info and 'columns' in describe_info:
                    print(f"✅ 表 {first_table} 的结构:")
                    for col in describe_info['columns'][:3]:  # 只显示前3列
                        print(f"  📝 {col['Field']}: {col['Type']} {'(主键)' if col['Key'] == 'PRI' else ''}")
                    if len(describe_info['columns']) > 3:
                        print(f"  ... 还有 {len(describe_info['columns']) - 3} 个字段")
                else:
                    print("❌ 无法获取表结构")
            else:
                print("❌ 表结构查询失败")
        else:
            print("\n❌ 没有可用的表进行结构演示")
        
        print("\n🎉 演示完成！")
        print("\n💡 使用提示:")
        print("1. 在Trae AI中，你可以直接使用这些MCP工具")
        print("2. 工具会自动处理数据库连接和查询")
        print("3. 所有操作都有详细的日志记录")
        
        # 关闭服务器
        await mcp_server.stop()
        
    except Exception as e:
        print(f"❌ 演示过程中出现错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(demo_mcp_tools())